#include "tracer.h"

#include <omp.h>
#include "material/light_material.h"
#include "material/object_material.h"
#include "material/color_material.h"
#include "material/texture_material.h"
#include "local_illumination.h"
#include "common/random_generator.h"

namespace Sirat {

namespace BDPT {

Common::Random_Generator random(-1, 1);

Common::Color get_specular(Object::Object* _object) {
	Material::Object_Material* material = _object->get_material();
	Common::Color specular = Common::Color(1.0,1.0,1.0);
	switch(material->get_type()) {
		case Material::Color_Material_Type: {
			specular = static_cast<Material::Color_Material*>(material)->get_specular();
			break;
		}
		default:
			break;
	}

	return specular;
}

Common::Color get_diffuse(Object::Object* _object) {
	Material::Object_Material* material = _object->get_material();
	Common::Color diffuse = Common::Color(1.0,1.0,1.0);
	switch(material->get_type()) {
		case Material::Color_Material_Type: {
			diffuse = static_cast<Material::Color_Material*>(material)->get_diffuse();
			break;
		}
		default:
			break;
	}

	return diffuse;
}

Tracer::Tracer(Common::Scene* _scene, Common::Camera* _camera, const unsigned int _path_number, const unsigned int _eye_path_depth, const unsigned int _light_path_depth) : scene(_scene), camera(_camera), path_number(_path_number), eye_path_depth(_eye_path_depth), light_path_depth(_light_path_depth) {
	srand(time(NULL));
	image = 0;
	#ifdef _OPENMP
	parallel = false;
	threads_number = 0;
	#endif
}

Tracer::Tracer(const Tracer& _other) {
	scene = _other.scene;
	camera = _other.camera;
	image = _other.image;
	path_number = _other.path_number;
	eye_path_depth = _other.eye_path_depth;
	light_path_depth = _other.light_path_depth;

	#ifdef _OPENMP
	parallel = _other.parallel;
	threads_number = _other.threads_number;
	#endif
}

void Tracer::render() {
	Common::Grid* grid = camera->get_grid();
	unsigned short height = grid->get_height();
	unsigned short width = grid->get_width();

	#ifdef _OPENMP
	bool parallel_test = parallel;
	#endif

	int x;
	#ifdef _OPENMP
	if(threads_number > 0) {
		omp_set_num_threads(threads_number);
	}
	#pragma omp parallel for default(shared) private(x) if(parallel_test) schedule(dynamic,10)
	#endif
	for(x = 0;x < width;x++) {
		for(unsigned int y =0;y < height;y++) {
			Math::Ray ray = camera->generate_ray(x,y);

			tuple<Object::Object*,Math::Vector> intersection = scene->nearest_object(ray);
			Object::Object* object = intersection.get<0>();
			Common::Color color;
			if (object != 0) {	 // 0 means no intersection 
				Math::Vector point = intersection.get<1>();
				color = calculate_color(object,ray,point);
			}
			#ifdef _OPENMP
				#pragma omp critical (set_pixel)
			#endif
			{
				if(image != 0) image->set_pixel(x,y,color);
			}
		}
	}
}

Common::Color Tracer::calculate_color(Object::Object* _object, Math::Ray _ray, Math::Vector _point) {
	Common::Color result;

	for(unsigned int count = 0; count < path_number; count++) {
		Path eye_path, light_path;
		Light::Light* random_light;

		unsigned int light_index = rand()%(scene->get_lights_size());
		random_light = scene->get_light(light_index);

		construct_path(eye_path, _object, _point, _ray, eye_path_depth);
		construct_path(light_path, random_light, light_path_depth);

		Common::Color estimate = calculate_estimate(eye_path, light_path);

		result = result + estimate;
	}

	result = result * (1.0/path_number);

	return result;
}

Math::Ray Tracer::get_random_ray(Math::Vector _point, Math::Vector normal) {
	Math::Vector direction;
	do {
		direction.x = random.generate();
		direction.y = random.generate();
		direction.z = random.generate();
	} while(direction.length() > 1.0 || direction.dot_product(normal) < 0.0);

	return Math::Ray(_point, direction.normalize());
}
	
void Tracer::construct_path(Path& _path, Object::Object* _object, Math::Vector _point, Math::Ray _ray, unsigned int _depth) {
	_path.push_back( make_tuple(_object, _point, Local_Illumination::calculate_color(_object, _ray, _point, scene) ));
	if(_depth == 0) return;

	Math::Vector normal = _object->get_normal(_point);
	Math::Ray ray = get_random_ray(_point, normal);
	tuple<Object::Object*,Math::Vector> intersection = scene->nearest_object(ray);
	if(intersection.get<0>() == 0) return;
	construct_path(_path, intersection.get<0>(), intersection.get<1>(), ray, _depth - 1);
}

void Tracer::construct_path(Path& _path, Light::Light* _light, unsigned int _depth) {
	if(_depth == 0) return;
	Math::Ray* rays;
	
	rays = _light->get_sample_directions(random, 1);
	tuple<Object::Object*,Math::Vector> intersection = scene->nearest_object(rays[0]);
	
	if(intersection.get<0>() != 0) construct_path(_path, intersection.get<0>(), intersection.get<1>(), rays[0], _depth - 1);

	delete[] rays;
}

Common::Color Tracer::calculate_estimate(Path eye_path, Path light_path) {
	if(eye_path.size() == 0) return Common::Color(0.0,0.0,0.0);
	Common::Color result;
	Common::Color weight_sum;
	
	for(unsigned int size_eye = 1; size_eye <= eye_path.size(); size_eye++) {
		for(unsigned int size_light = 0; size_light <= light_path.size(); size_light++) {
			Path result_path;

			for(unsigned int i = 0; i < size_eye; i++) result_path.push_back(eye_path[i]);
			for(unsigned int i = 0; i < size_light; i++) {
				if(i == 0) {
					Math::Vector eye_point = eye_path[size_eye - 1].get<1>();
					Math::Vector light_point = light_path[0].get<1>();

					double distance = eye_point.distance(light_point);
					Math::Vector direction = (light_point - eye_point).normalize();

					if(scene->object_occlusion(Math::Ray(eye_point, direction), distance)) break;
				}

				result_path.push_back(light_path[i]);
			}

			Common::Color diffuse = get_diffuse(result_path[0].get<0>());
			Common::Color specular = get_specular(result_path[0].get<0>());
			Common::Color sum = diffuse + specular;
			Common::Color path_weight;
			{
				double r = sum.red;
				double g = sum.green;
				double b = sum.blue;

				path_weight = Common::Color(1.0 - 1.0/(r + 0.000001), 1.0 - 1.0/(g + 0.000001), 1.0 - 1.0/(b + 0.000001));
			}
			
			weight_sum = weight_sum + path_weight;
			result = result + path_estimate(result_path)*path_weight;
		}
	}

	weight_sum.red = 1.0/(weight_sum.red + 0.0000001);
	weight_sum.green = 1.0/(weight_sum.green + 0.0000001);
	weight_sum.blue = 1.0/(weight_sum.blue + 0.0000001);

	result = result*weight_sum;

	return result;
}

Common::Color Tracer::path_estimate(Path _path, unsigned int _index) {
	if(_path.size() == 0) return Common::Color(0.0,0.0,0.0);

	Object::Object* object = _path[_index].get<0>();
	Common::Color diffuse = get_diffuse(object);

	if(_index == _path.size() - 1) return _path[_index].get<2>()*diffuse;

	Common::Color result;
	Common::Color next_result = path_estimate(_path, _index + 1);
	
	Math::Vector point = _path[_index].get<1>();
	Math::Vector next_point = _path[_index + 1].get<1>();
	
	Common::Color local = _path[_index].get<2>();

	Math::Vector normal = object->get_normal(point);
	double normal_factor = normal.dot_product((next_point - point).normalize());
	result = local + next_result * normal_factor * diffuse;
	return result;
}


void Tracer::set_output(Common::Image* _image) {
	image = _image;
}

#ifdef _OPENMP
void Tracer::set_parallel(const bool _parallel) {
	parallel = _parallel;
}

void Tracer::set_threads(const unsigned int _threads) {
	threads_number = _threads;
}
#endif

}

}
